import numpy as np
class KMeans:
    def __init__(self, k=3, max_iters=100, tol=1e-4):
        self.k = k
        self.max_iters = max_iters
        self.tol = tol
    def fit(self, X):
        np.random.seed(42)
        # Step 1：随机选择 k 个点作为初始聚类中心
        random_idx = np.random.choice(len(X), self.k, replace=False)
        self.centroids = X[random_idx]
        for i in range(self.max_iters):
            # Step 2：分配样本到最近的聚类中心
            labels = self._assign_clusters(X)
            # Step 3：计算新的聚类中心
            new_centroids = np.array([X[labels == j].mean(axis=0) for j in range(self.k)])
            # Step 4：判断收敛（中心点移动是否足够小）
            if np.all(np.abs(new_centroids - self.centroids) < self.tol):
                break
            self.centroids = new_centroids
        self.labels_ = labels
    def _assign_clusters(self, X):
        distances = np.linalg.norm(X[:, np.newaxis] - self.centroids, axis=2)
        return np.argmin(distances, axis=1)
    def predict(self, X):
        return self._assign_clusters(X)
if __name__ == "__main__":
    # 构造两类样本
    X1 = np.random.randn(100, 2) + np.array([2, 2])
    X2 = np.random.randn(100, 2) + np.array([-2, -2])
    X3 = np.random.randn(100, 2) + np.array([5, -3])
    X = np.vstack((X1, X2, X3))
    # 聚类
    kmeans = KMeans(k=3)
    kmeans.fit(X)
    print("聚类中心：")
    print(kmeans.centroids)
    print("前 10 个样本的类别：")
    print(kmeans.labels_[:10])